import requests
import os
# 定义服务器的URL和端口
SERVER_URL = 'http://localhost:5000'


# 定义下载预测结果的函数
def download_prediction(output_dir):
    prediction_url = f"{SERVER_URL}/get_prediction/{output_dir}"
    response = requests.get(prediction_url, stream=True)

    # 检查响应状态码
    if response.status_code == 200:
        # 将文件保存到本地
        print(prediction_url)
        file_name = os.path.basename(prediction_url).split('/')[-1]
        file_name = f'{file_name}.wav'
        with open(file_name, 'wb') as file:
            for chunk in response.iter_content(chunk_size=1024):
                if chunk:  # 过滤掉保持连接的新块
                    file.write(chunk)
        print(f"预测文件 '{file_name}' 下载成功。")
    elif response.status_code == 404:
        print("错误：未找到预测结果。")
    else:
        print(f"错误：下载预测文件失败，状态码：{response.status_code}")


# 主函数，用于运行客户端
def main():
    # 假设我们有一个输出目录ID
    output_dir_id = '123'

    # 下载预测结果文件
    download_prediction(output_dir_id)


if __name__ == "__main__":
    main()